package corenlp;

import edu.stanford.nlp.coref.CorefCoreAnnotations;
import edu.stanford.nlp.coref.data.CorefChain;
import edu.stanford.nlp.coref.data.Mention;
import edu.stanford.nlp.ling.*;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.pipeline.StanfordCoreNLP;
import edu.stanford.nlp.semgraph.SemanticGraph;
import edu.stanford.nlp.semgraph.SemanticGraphCoreAnnotations;
import edu.stanford.nlp.trees.Constituent;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.trees.TreeCoreAnnotations;
import edu.stanford.nlp.util.CoreMap;
import edu.stanford.nlp.util.IntPair;
import scala.tools.cmd.gen.AnyVals;

import java.util.*;

/**
 * Created by common on 2017/4/6.
 */

public class BasicPipelineExample {


    public static void interprete_documents(Annotation document){
        // these are all the sentences in this document
        // a CoreMap is essentially a Map that uses class objects as keys and has values with custom types
        List<CoreMap> sentences = document.get(CoreAnnotations.SentencesAnnotation.class);

        for(CoreMap sentence: sentences) {
            // traversing the words in the current sentence
            // a CoreLabel is a CoreMap with additional token-specific methods
            for (CoreLabel token: sentence.get(CoreAnnotations.TokensAnnotation.class)) {
                // this is the text of the token
                String word = token.get(CoreAnnotations.TextAnnotation.class);
                // this is the POS tag of the token
                String pos = token.get(CoreAnnotations.PartOfSpeechAnnotation.class);
                // this is the NER label of the token
                String ne = token.get(CoreAnnotations.NamedEntityTagAnnotation.class);

                System.out.println("word=" + word + " pos=" + pos + " ne=" + ne);
            }



            // this is the parse tree of the current sentence
            Tree tree = sentence.get(TreeCoreAnnotations.TreeAnnotation.class);

            // this is the Stanford dependency graph of the current sentence
            SemanticGraph dependencies = sentence.get(SemanticGraphCoreAnnotations.CollapsedCCProcessedDependenciesAnnotation.class);
        }

        // This is the coreference link graph
        // Each chain stores a set of mentions that link to each other,
        // along with a method for getting the most representative mention
        // Both sentence and token offsets start at 1!
        Map<Integer, CorefChain> graph =
                document.get(CorefCoreAnnotations.CorefChainAnnotation.class);
    }





    public static Annotation annotate_pipeline() {

        // creates a StanfordCoreNLP object, with POS tagging, lemmatization, NER, parsing, and coreference resolution
        Properties props = new Properties();

        // coref 耗内存，dcoref 省内存
        //props.setProperty("annotators", "tokenize, ssplit, pos, lemma, ner, parse, dcoref");
        props.setProperty("annotators", "tokenize,ssplit,pos,lemma,ner,parse,mention,coref");
        props.setProperty("coref.algorithm", "neural");
        StanfordCoreNLP pipeline = new StanfordCoreNLP(props);

        // read some text in the text variable
        String text = "Barack Obama is the president of USA. Obama was born in Honolulu. He graduated from Harvard Law School, where he met his wife.";
        //String text = "The Rhubarb have medicinal roots and edible stalks, however its leaves are toxic.";

        // create an empty Annotation just with the given text
        Annotation document = new Annotation(text);

        // run all Annotators on this text
        pipeline.annotate(document);

        System.out.println("\n\n\n\n\n");

        pipeline.prettyPrint(document, System.out);

        //pipeline.

        return document;

    }



    public static void interprete_mention(Annotation document){
        for (CoreMap sentence : document.get(CoreAnnotations.SentencesAnnotation.class)) {
            System.out.println("\n\n---");
            System.out.println("mentions");
            for (Mention m : sentence.get(CorefCoreAnnotations.CorefMentionsAnnotation.class)) {
                System.out.println("\t" + m);
                int begin = m.headIndex;
                int end = m.endIndex;
                System.out.println("[" + begin + " , " + end + "]");
            }
        }
    }




    public static void interprete_coref(Annotation document){
        System.out.println("coref chains");
        for ( CorefChain cc : document.get(CorefCoreAnnotations.CorefChainAnnotation.class).values() ) {
            System.out.println("\t" + cc);

        }

        // This is the coreference link graph
        // Each chain stores a set of mentions that link to each other,
        // along with a method for getting the most representative mention
        // Both sentence and token offsets start at 1!
        Map<Integer, CorefChain> graph =
                document.get(CorefCoreAnnotations.CorefChainAnnotation.class);
        System.out.println("Print The Graph.................");
        for(Integer key: graph.keySet()){
            System.out.println("\n\n\nkey = " + key);
            CorefChain cc = graph.get(key);
            System.out.println(cc);


            Map<IntPair, Set<CorefChain.CorefMention>> mentionMap = cc.getMentionMap();
            for(IntPair pair: mentionMap.keySet()){
                int source = pair.getSource();
                int target = pair.getTarget();
                Set<CorefChain.CorefMention> mentionSet = mentionMap.get(pair);
                System.out.println("pair = " + pair.toString() + " , source = " + source + " , target = " + target + " , size = " + mentionSet.size());

                for(CorefChain.CorefMention mention: mentionSet){
                    System.out.println("   ----->  " + mention);
                }
            }

            CorefChain.CorefMention mention = cc.getRepresentativeMention();
            System.out.println("mention = " + mention);
            //cc.deleteMention(mention);
            //System.out.println(cc);

        }
    }




    public static String resolve_coref(Annotation doc){



        Map<Integer, CorefChain> corefs = doc.get(CorefCoreAnnotations.CorefChainAnnotation.class);
        List<CoreMap> sentences = doc.get(CoreAnnotations.SentencesAnnotation.class);


        List<String> resolved = new ArrayList<String>();

        for (CoreMap sentence : sentences) {

            List<CoreLabel> tokens = sentence.get(CoreAnnotations.TokensAnnotation.class);

            for (CoreLabel token : tokens) {

                System.out.println();
                System.out.println();
                System.out.println("-----------------------------------------------------------------");

                Integer corefClustId= token.get(CorefCoreAnnotations.CorefClusterIdAnnotation.class);
                System.out.println(token.word() +  " --> corefClusterID = " + corefClustId);


                CorefChain chain = corefs.get(corefClustId);
                System.out.println("matched chain = " + chain);


                if(chain==null){
                    resolved.add(token.word());
                }else{

                    CorefChain.CorefMention reprMent = chain.getRepresentativeMention();
                    System.out.println(reprMent);

                    int sentINdx = reprMent.sentNum -1;
                    CoreMap corefSentence = sentences.get(sentINdx);
                    List<CoreLabel> corefSentenceTokens = corefSentence.get(CoreAnnotations.TokensAnnotation.class);


                    String newwords = "";

                    System.out.println("token.sentIndex = " + token.sentIndex() + " , rep.sentNum = " + sentINdx);
                    System.out.println("token.index = " + token.index() + " , rep.start = " + reprMent.startIndex + " , rep.end = " + reprMent.endIndex);
                    if(token.sentIndex() != sentINdx || token.index() < reprMent.startIndex || token.index() > reprMent.endIndex ){
                        for(int i = reprMent.startIndex; i<reprMent.endIndex; i++){
                            CoreLabel matchedLabel = corefSentenceTokens.get(i-1); //resolved.add(tokens.get(i).word());
                            resolved.add(matchedLabel.word());

                            newwords+=matchedLabel.word()+" ";
                        }
                    }
                    else{
                        resolved.add(token.word());
                    }


                    System.out.println("converting " + token.word() + " to " + newwords);
                }

                System.out.println("-----------------------------------------------------------------");

            }

        }


        String resolvedStr ="";
        System.out.println();
        for (String str : resolved) {
            resolvedStr+=str+" ";
        }
        System.out.println(resolvedStr);

        return resolvedStr;


    }



    public static String traversal_parsetree(Annotation document){

        Map<Integer, CorefChain> corefs = document.get(CorefCoreAnnotations.CorefChainAnnotation.class);


        // these are all the sentences in this document
        // a CoreMap is essentially a Map that uses class objects as keys and has values with custom types
        List<CoreMap> sentences = document.get(CoreAnnotations.SentencesAnnotation.class);

        for(CoreMap sentence: sentences) {
            System.out.println("\n\nsentence------");
            // traversing the words in the current sentence
            // a CoreLabel is a CoreMap with additional token-specific methods
            List<CoreLabel> tokens = sentence.get(CoreAnnotations.TokensAnnotation.class);
            for (CoreLabel token : tokens) {
                // this is the text of the token
                String word = token.get(CoreAnnotations.TextAnnotation.class);
                // this is the POS tag of the token
                String pos = token.get(CoreAnnotations.PartOfSpeechAnnotation.class);
                // this is the NER label of the token
                String ne = token.get(CoreAnnotations.NamedEntityTagAnnotation.class);

                Integer index = token.get(CoreAnnotations.IndexAnnotation.class);

                System.out.println("word=" + word + " pos=" + pos + " ne=" + ne + " index=" + index);
            }


            // this is the parse tree of the current sentence
            Tree tree = sentence.get(TreeCoreAnnotations.TreeAnnotation.class);
            tree.pennPrint();
            tree.indentedXMLPrint();
            tree.indentedListPrint();
            System.out.println("\n\ntree = \n" + tree.toString());
            Set<Constituent> constituentSet = tree.constituents();
            for(Constituent constituent: constituentSet){
                System.out.println(constituent.toString());
            }

            System.out.println("\nTraversal Tree...........");
            traversal_recursive(tree, corefs, tokens,"");



        }





        return "";
    }


    public static void traversal_recursive(Tree tree, Map<Integer, CorefChain> corefs, List<CoreLabel> tokens, String indent){
        Label label = tree.label();
        System.out.println(indent + label);

        // http://stackoverflow.com/questions/16502378/stanford-corenlp-get-characteroffset-annotation-from-parse-tree
        HasOffset hasOffset  = (HasOffset) label;
        int begin = hasOffset.beginPosition();
        int end = hasOffset.endPosition();
        System.out.println("[" + begin + "," + end + "]");

        //HasIndex has_index = (HasIndex) label;
        //System.out.println("[has_index= " + has_index + "]");

        //HasCategory hasCategory = (HasCategory)label;
        //System.out.println("[hasCategory= " + hasCategory + "]");

        //HasContext hasContext = (HasContext)label;
        //hasContext.
        //System.out.println("[hasContext= " + hasContext + "]");

        //HasTag hasTag = (HasTag)label;
        //hasTag.
        //System.out.println("[hasTag= " + hasTag + "]");

        //HasWord hasWord = (HasWord)label;
        //hasWord.
        //System.out.println("[hasWord= " + hasWord + "]");

        //HasLemma hasLemma = (HasLemma)label;
        //System.out.println("[hasLemma= " + hasLemma + "]");



        //if(begin!=-1 && end!=-1){
        if(tree.isLeaf()){
            String[] labelArray = label.toString().split("-");
            String indexStr = labelArray[labelArray.length-1];
            System.out.println("[indexStr= " + indexStr + " ]");
            Integer index = Integer.parseInt(indexStr);
            System.out.println("[index= " + index + " ]" );

            CoreLabel matchedToken = tokens.get(index-1);
            Integer matchedIndex = matchedToken.get(CoreAnnotations.IndexAnnotation.class);
            System.out.println("[matchedIndex= " + matchedIndex + " ]");

            Integer corefClustId= matchedToken.get(CorefCoreAnnotations.CorefClusterIdAnnotation.class);
            System.out.println("[" + matchedToken.word() +  " --> corefClusterID = " + corefClustId + " ]");

            CorefChain chain = corefs.get(corefClustId);
            System.out.println("[matched chain = " + chain + "]");

        }


        //IntPair span = tree.getSpan();
        //System.out.println(indent + "span = (" + span.getSource() + "," + span.getTarget() + ")");

        List<Tree> childTrees = tree.getChildrenAsList();
        for(Tree child : childTrees){
            traversal_recursive(child, corefs, tokens,indent + "  ");
        }
    }


    public static void main(String[] args) {

        Annotation document = annotate_pipeline();

        //interprete_mention(document);

        interprete_coref(document);

        //resolve_coref(document);

        //traversal_parsetree(document);

        //interprete_documents( document );

    }

}